from random import random
from textwrap import shorten
import PIL
import traceback
from matplotlib import image
from matplotlib.pyplot import sca
from numpy import add_newdoc, mask_indices
from image_angle_correct import *
from hsv_test import hsv_split, show,bigger_box,hsv_split_pre
from model import *
import sys
sys.path.append("../")
from util.keys_prov import *
from PIL import Image, ImageDraw, ImageFont
'''
    识别流程:
    1.hsv掩膜
    2.识别文字框，旋转纠正角度
    3.识别，若无，则旋转180度，再识别，若仍无结果，接连3次旋转60度并识别
'''

# TEST 
# source_image_path = "/home/meroke/code/detection/result/error/"  # test
# result_image_save_path = "/home/meroke/code/detection/result/test/result/" # test
# error_image_save_path ="/home/meroke/code/detection/result/test/error/" # test 

# MAIN
# 训练集路径
source_image_path = "./image/test/" # main
# 识别结果图片保存路径
result_image_save_path = "./image/result/" # main
# 未识别地区的图片保存路径
error_image_save_path ="./image/error/" # mian


# 地区识别结果保存路径
word_detect_path = "log/word_detect.txt"
# 识别结果细节保存路径
detect_detail_path = "log/detect_detail.txt"

# 检测正确数量 计数器
detect_count =0
# 所有图片数量 计数器
image_count = 0

'''
description:  在图片上绘制中文信息
param {*} img  绘制图片
param {*} text  要绘制的中文信息
param {*} position 绘制的像素点位置
param {*} textColor 绘制的文字颜色
param {*} textSize 绘制的文字大小
return {*} 返回BGR格式的，绘制有中文信息的图片
'''
def cv2AddChineseText(img, text, position, textColor=(0, 255, 0), textSize=30):
    if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    if text is "":
        text = "安徽"
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img)
    # 字体的格式
    fontStyle = ImageFont.truetype(
        "simsun.ttc", textSize, encoding="utf-8")
    # 绘制文本
    draw.text(position, text, textColor, font=fontStyle)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)


'''
description: 字符串匹配省份
param {*} detected_word 检测到的字符串 列表
return {*} 返回匹配的省份名称， 无结果则返回空字符串
'''
def searcch_provience(detected_word):
    result = ""
    flag = False
    for key in province:
        if flag:
            break
        for single in detected_word[1]:
            if single.find(province[key])>-1 or single.find(key)>-1:
                result=province[key]
                # print(true_result)
                flag=True
                break
    return result

'''
description: 记录识别到的省份 及 检测出省份的图片名称
param {*} name 检测出省份的图片名称
param {*} result 省份名称
'''
def record_result(name,result):
    global word_detect_path
    if not os.path.exists("log"):
        os.makedirs("log")
    with open(word_detect_path, 'a') as f:
        f.write(result)  #文件的写操作
        f.write(" " + name + ".jpg\r")

def record_error(name,result):
    global word_detect_path
    if not os.path.exists("log"):
        os.makedirs("log")
    with open(word_detect_path, 'a') as f:
        f.write(result)  #文件的写操作
        f.write(" " + name + ".jpg\r")
'''
description:  记录识别到的所有字符串 及 对应的图片名称
param {*} name 图片名称 
param {*} detail 识别到的所有字符串
'''
def record_detail(name,detail):
    global detect_detail_path
    with open(detect_detail_path, 'a') as f:
        for i in range(len(detail)):
            f.write(detail[i] +" ")  #文件的写操作
        f.write(" " + name + ".jpg\r")

'''
description:  保存省份未匹配成功的图片（缩小尺寸）
param {*} name 图片名称（路径）
param {*} img 图片
'''
def error_img_save(name,img):
    global error_image_save_path
    try:
        img.save(os.path.join(error_image_save_path , name + ".jpg"))
    except:
        cv2.imwrite(os.path.join(error_image_save_path , name),img)


'''
description: 匹配文字框与轮廓框，获取包含文字框的轮廓
param cnt 轮廓框
param text_box 文本框
return {*} 返回包含文字框的轮廓，否则为None
'''
def get_bbox(cnt,text_box):
    for i in cnt:
        # print(i)
        if (text_box[0][0] > i[0][0]-30 and text_box[0][1] > i[0][1]-30 and\
            text_box[1][0] < i[1][0]+30 and text_box[1][1] > i[1][1]-30):
            # print("outpout ",i)
            return i
    return None

'''
description: 获取轮廓中心坐标
param box 轮廓
return {*} 中心坐标
'''
def get_bbox_center(box):
    x = (box[1][0] - box[0][0])/2 +  box[0][0]
    y = (box[3][1] - box[0][1])/2 +  box[0][1]

    return (int(x),int(y))

'''
description: 获取识别结果，匹配省份信息，返回省份
param img 最终处理后的图片
param name 图片名称
param origin_img 原始图片
param short_size 图片短边长度
return {*}
'''
def get_result(img,name,origin_img,short_size):
    PIL_image = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
    # time.sleep(1) 
    detected_word = OcrHandle().text_predict(PIL_image,short_size)
    result = searcch_provience(detected_word)
    # print("result:",result)
    if result is "":
        PIL_image = PIL_image.transpose(Image.ROTATE_180)
        #show("04-2Rotatede",cv2.cvtColor(np.asarray(PIL_image), cv2.COLOR_RGB2BGR))
        detected_word = OcrHandle().text_predict(PIL_image,short_size)
        result = searcch_provience(detected_word)
        if result is "":
            for i in range(4):
                PIL_image = PIL_image.transpose(Image.ROTATE_90)
                detected_word = OcrHandle().text_predict(PIL_image,short_size)
                result = searcch_provience(detected_word)
                if result is not "":
                    break
            error_img_save(name,origin_img)

    return result,detected_word

def resize_image(img,scale=10):
    h,w = img.shape[0],img.shape[1]
    if(h>1024 or w > 1024):
        scale = h /1024
        h = int(h/scale)
        w = int(w/scale)
        img= cv2.resize(img,(h,w))
    return img,h,w

'''
description: 文字识别 主要流程函数
param path 图像路径
param name 图像名称
return {*}
'''
def word_detect_test(path,name):
    global detect_count, image_count, result_image_save_path, error_image_save_path
    print("{} START".format(name))
    img = cv2.imread(path)
    img,h,w =  resize_image(img)
    short_size = h if h<w else w
    print(w,h,short_size)
    origin = img.copy()
    #show("01origin",img)

    # 第一次蓝色分割
    img_pre = hsv_split_pre(origin,h,w)
    #show("02hsv_split",img_pre)
    # 第二次白色分割
    cnt, img = hsv_split(img_pre,h,w)
    #show("02-2hsv_split",img)

    # 返回识别框的坐标点列表
    box_list, score_list = text_handle.process(img,short_size)
    if len(box_list) < 1:
        print("detect fail {}".format(name))
        error_img_save(name,origin)
        return
    # 获取最大识别框下标
    max_index = max_box_index(box_list,img)

    # 绘制展示最大识别框
    img = draw_max_bbox(img, box_list[max_index])
    # 根据固定下标0 1 两点 计算图片应该旋转的角度
    angle = get_angle(img,box_list[max_index])

    # 图像角度矫正
    img = adjust_angle(img,angle)
    #show("03adjust_angle",img)

    # 图像锐化
    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], np.float32) #定义一个核
    img = cv2.filter2D(img, -1, kernel=kernel) 
    #show("04sharp_img",img)

    # 文字识别 获取最终识别结果
    result,detected_word  = get_result(img,name,origin,short_size)
    if result is not "":
        detect_count +=1
    # 如果未成功识别，自动标注为安徽
    else:
        result = "安徽"
    image_count +=1

    print("detected:{}\r".format(result))

    # 在图像上绘制最终显示的省份
    CV_image = cv2AddChineseText(origin,result,(50,100),(255,0,0),50)
    main_bbox =  get_bbox(cnt,box_list[max_index])
    print("main_bbox: " ,main_bbox)
    if main_bbox is None:
        # record_error(name,"main_box maybe is None")
        traceback.print_exc(file=open("log/error.log",'a'))
        main_bbox = np.array(bigger_box(box_list[max_index],100))
        
    # 绘制最终显示文本框
    cv2.waitKey(0)
    if(main_bbox.all() != None):
        try:
            CV_image = cv2.drawContours(CV_image,[main_bbox],-1,(0,0,255),2)
        except:
            print("main_bbox error")
        #show("05-1bboximage",CV_image)
    center = get_bbox_center(main_bbox)
    cv2.putText(CV_image,"["+str(center[0])+" "+str(center[1])+"]",center,
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
    #show("05result",CV_image)

    # 记录结果
    name = name.split('.')[0]
    cv2.imwrite(result_image_save_path +"{}_del.jpg".format(name), CV_image)
    # 记录任务要求信息
    record_result(name,result+",("+str(center[0])+" "+str(center[1])+")")
    # 记录详细信息
    record_detail(name,detected_word[1])


'''
description: 清空所有记录表，便于记录本次识别数据
return {*}
'''
def prepare():
    global word_detect_path
    if os.path.exists(word_detect_path):
        os.remove(word_detect_path)
    if os.path.exists(detect_detail_path):
        os.remove(detect_detail_path)
    if os.path.exists(result_image_save_path):
        for i in os.listdir(result_image_save_path):
            os.remove(os.path.join(result_image_save_path,i))
    if os.path.exists(error_image_save_path):
        for i in os.listdir(error_image_save_path):
            os.remove(os.path.join(error_image_save_path,i))
    print("prepare done!")

if __name__ == "__main__":
    
    prepare()
    start_time = time.time()
    file  = os.listdir(source_image_path)
    random_list =[]
    for i in os.listdir(source_image_path):
        if "jpg" in i:
            if "del" in i:
                continue
            # hsv_split(source_image_path + i,i)
            word_detect_test(source_image_path + i,i)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
            print(i + " done!")

    end_time  = time.time()
    print("cost time:{} s".format(end_time - start_time))
    print("cost time:{} m".format((end_time - start_time)/60))

    print("sum:{} detected:{}, rate:{}".format(image_count, detect_count, detect_count/image_count) )